Qwen3.5-397B RL trainer on 128 v7x: the tuned recipe turns on the fp8 fake-quant hoist, fused SwiGLU backward, attention savings, the 2048 wi forward tile and the GDN context-parallel m scan, and pins the MaxText that defines them - #2948

…joins TRAINER_TUNED_FLAGS

TRAINER_TUNED_FLAGS=true now also passes attention's load-balance reorder of
the positions and segment ids, kept out of the remat replay: the one MaxText
change of the measured cell that the block did not pass.

In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains
context_parallel_attention_load_balance_hoist_ids=true, before the block's
custom_mesh_and_rule=cp-as-ep-dense-no-cp: remat_policy=custom saves each
attention layer's load-balance-ordered positions and segment ids. They depend
only on the batch, so jax.lax.scan keeps one copy outside the layer loop, and
the remat replay of every attention layer reads it instead of reordering them
again. The values are the same. Compiled for TPU v7x, XLA already moves the
reorder's collective-permutes out of the loop without the key; the key moves
out the rest of the reorder and the work on its result (per the MaxText
commit: an AOT compile of the bench's 128-chip trainer program at an earlier
cell has 47 fewer instructions in its while loops with the change, the same
collective-permutes and the same memory). MaxText refuses the key without
context_parallel_attention_load_balance=true, an even context parallelism of at
least 2, remat_policy=custom and scan_layers=true. This recipe's trainer has
them at 32x2x4 and at the 256-chip meshes 32x2x8 and 64x1x8 (context 8), so
the key goes on every trainer mesh, like the block's other keys.

The key is defined in the MaxText change on top of AI-Hypercomputer/maxtext#5721
(branch parambole/rl-4600-final-maxtext), which the previous commit pins;
atwigg/mlperf at bbd0f6f10 does not define it. TRAINER_TUNED_FLAGS=false and
the other 397B v7x recipes do not pass it.

Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x,
mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches
per step, mean of steps 1 to 4, no token dropped in any arm), alone, in three
comparisons against controls in other pods whose arm commands differ only by
the change's switch: 99.504 -> 99.267 s per step (o275x-gdn-ctl ->
o275x-s3-x17h), 99.248 -> 99.267 (o275x-mf-ctl -> o275x-s3-x17h) and
91.7513 -> 91.8665 (o275x-stk3-l3-prof -> o275x-s5c-l3-prof), i.e. -0.237,
+0.019 and +0.115 s per step. One unchanged configuration spans 0.255 s per
step over four pods (o275x-hd), so no effect is resolved. Together with the
next commit's two keys and the mesh rule, see the next commit (F2).

This is 11c096bb0 (branch parambole/rl-trainer-4600-next of google/tunix), as
ported onto #2884, replayed onto #2884's head b16ee52. Two hunks
conflicted, both comments that #2884 had just rewritten for its new pin: the
recipe's sentence on where the tuned keys are defined and the test's comment
on _BENCH_MAXTEXT_KEYS. Each keeps #2884's wording and adds the key. The
test's comment on _MAXTEXT_DEFAULTS now says the change is on top of #5721.

Tests (mlperf_submission_launcher_test.py): _MAXTEXT_DEFAULTS has the key's
default (false), the exact-set test expects it among _TUNED_KEYS, and the
switch-off and other-recipe tests check that it is not passed
(_BENCH_MAXTEXT_KEYS).
In the bench image (dbdcf5f8dcba; CPU, no network) 48 of 48 pass at this
commit. Rendered, the submission launch differs from the parent's only by this
key in the trainer's MaxText flags, at 32x2x4, 32x2x8 and 64x1x8; with
TRAINER_TUNED_FLAGS=false, and the other 397B v7x recipes at their defaults,
the renders are byte-identical.

…her_once join TRAINER_TUNED_FLAGS

TRAINER_TUNED_FLAGS=true now also turns on the two remaining keys of the
MaxText change on top of AI-Hypercomputer/maxtext#5721 (branch
parambole/rl-4600-final-maxtext, pinned two commits ago): the SwiGLU backward
fused into the down projection's input-gradient GMM, and the query kernel
gathered once per attention layer in the backward.

In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains:

- moe_wo_gmm_fused_dswiglu=true, on every trainer mesh: the input-gradient GMM
  of the MoE down projection turns its float32 accumulator into d[gate | up]
  in its epilogue (with moe_router_grad_from_wo_input, also the routing
  weight's gradient) instead of writing the activation's gradient in bfloat16
  for separate passes to read back. The gradients of wi and of everything
  before it change by bfloat16 rounding. MaxText refuses the key unless that
  GMM has one n tile (wo_tile_dlhs_mlp_dim = the MoE MLP dim, 1024 in this
  recipe) and, with moe_wi_gmm_fused_swiglu, wi_tile_fwd_mlp_dim halves to
  that dim or to a multiple of 128 dividing it (this recipe's 1024 gives two n
  tiles of 512 per half); it also needs explicit tiles, prefuse_moe_weights,
  use_tokamax_gmm, use_gmm_v2 and a remat policy that does not save the routed
  experts' gate or up half.
- query_proj_kernel_gather_once=true, at every trainer mesh except the
  256-chip 32x2x8 and 64x1x8: the attention query projection keeps the kernel
  it gathered in the remat replay for its input gradient, so the backward
  gathers the query kernel once per attention layer instead of twice. The
  values are the same. MaxText refuses the key unless the backward replays the
  query projection (remat_policy=custom with query_proj=remat). #2884's block
  sets query_proj=device at 32x2x8 and 64x1x8, where the backward does not
  replay the projection, so the key goes in a new default branch of that mesh
  case instead of on the tuned line: passed there, MaxText would refuse the
  256-chip trainer when it starts.

Measured on TPU together (F2): the bench cell of this block at 32x2x4 with
these two keys, context_parallel_attention_load_balance_hoist_ids and
cp-as-ep-dense-no-cp, on the MaxText of AI-Hypercomputer/maxtext#5758 with the
change on top of it (4ed82d793, whose src/maxtext is the pinned MaxText's
without #5721's three commits of Rohan Bierneni and two model YAMLs). Bench
(AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh
fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, padding first, 9
micro-batches per step, mean of steps 1 to 4, no token dropped), booking
r-7a4cc: 80.398 and 80.342 s per step in two pods (o275x-f2-l3-prof in
parambole-bf2-book, o275x-f2-l3-r in parambole-bf2r-book), 3,668 and 3,671
tokens/s/chip, regular micro-batch 8.748 and 8.749 s, against 82.373 s per
step, 3,580 tokens/s/chip and 8.966 s for the code of #5758 with the keys
before these four (o275x-prfc-f1 in parambole-bf1-book). The F2 pods' own
control arms did not finish, so the comparison is across pods: -1.974 and
-2.030 s per step. Numerics in both pods (nxf2: loss gates off, random
advantages, three steps): all six gates pass against that code under the same
protocol (nxprf-l3) and against the reference pair (nxref, nxref2); the loss
is 0.0304453575 in all three steps, nxprf-l3's to the ten digits logged, and
the gradient norm (gate N3) is within 1.97e-4 and 1.23e-4 relative of
nxprf-l3's (tolerance 2e-3). Each key's own measurement:
- moe_wo_gmm_fused_dswiglu, behind an environment switch on an overlay build
  with the keys of #2913's recipe, the hoist and padding first, without the
  fused log-prob head, against a control in another pod: 83.316 -> 82.148 s
  per step (o275x-c9pfxctl-l3 -> o275x-c9ypfctl-l3), 9.0845 -> 8.9545 s per
  regular micro-batch. On one TPU v7x chip its kernels at the cell's shapes
  match a float32 SwiGLU backward to 1.658e-3 relative L2
  (parambole-bck-2d4c2), and its numerics run passes all six gates
  (nxc9ypfckhv).
- query_proj_kernel_gather_once: not timed on TPU alone. An AOT compile of the
  cell with cp-as-ep gathers the query kernel once per attention layer in the
  backward instead of twice at the same device peak (104.46 GiB), and the
  MaxText change's test, executed on TPU v5p, gives bit-identical values and
  gradients with and without it.
Not run: with #5721's three commits (#5721's own content with the rule ran at
81.230 and 81.168 s per step in booking r-67724, o275x-rbd-v16 and
o275x-rbd-r-v16), at the 256-chip meshes, or end to end.

This is b7f7a6570 (branch parambole/rl-trainer-4600-next of google/tunix)
without its third value, custom_mesh_and_rule=cp-as-ep-dense-no-cp, which
#2884 already passes, ported onto #2884 and replayed onto its head b16ee52.
The port keeps #2884's structure: the header (#2884 rewrote it for the
256-chip meshes; it gains one line naming the rule under TRAINER_TUNED_FLAGS),
the end of the block's comment, the tuned line (#2884's, with
moe_wo_gmm_fused_dswiglu after the hoist key), the mesh case (the new default
branch) and, in the test, the defaults (the two keys join the MaxText-change
defaults ahead of #2884's per-tensor remat keys). On replay the same two
comments as in the previous commit conflicted (the recipe's sentence on
where the keys are defined and the test's comment on _BENCH_MAXTEXT_KEYS),
resolved the same way, now naming all three keys. The block's comment on
the whole cell replaces "has not run on TPU" with the F2 measurement above.

Tests, in the bench image (dbdcf5f8dcba; CPU, no network, the tree read-only,
stubs for kubectl, gcloud and gsutil first on PATH):
- mlperf_submission_launcher_test.py: 48 of 48 pass at this commit, at each of
  the two commits before it, and at #2884's b16ee52 and 8647d69.
- Five deliberately wrong copies of the recipe each fail exactly the tests
  predicted before the run: moe_wo_gmm_fused_dswiglu or
  context_parallel_attention_load_balance_hoist_ids dropped from the block
  (the 3 exact-set tests); the default branch not passing
  query_proj_kernel_gather_once (those 3, the default-cell test and the two
  256-chip mesh tests); query_proj_kernel_gather_once also on the tuned line
  (the 4 tests of the 256-chip meshes); moe_wo_gmm_fused_dswiglu passed only
  with TRAINER_TUNED_FLAGS=false (the 3 and the two switch-off tests). This
  test fails 6 tests on b16ee52's recipe, and b16ee52's test fails 5 on
  this recipe, both as predicted.
- Rendered, the submission launch differs from the parent's only in the
  trainer's MaxText flags: moe_wo_gmm_fused_dswiglu=true and
  query_proj_kernel_gather_once=true at 32x2x4, moe_wo_gmm_fused_dswiglu=true
  at 32x2x8 and 64x1x8. With TRAINER_TUNED_FLAGS=false, and the other 397B v7x
  recipes at their defaults, the renders are byte-identical.
- MaxText config check (the rendered trainer JobSet through
  run_trainer_node's argument parsing and build_maxtext_config to
  pyconfig.initialize, which runs the validators, with mocked devices): the
  pinned MaxText (94f11c09d) accepts the trainer at 32x2x4 (256 devices; all
  three keys on, query_proj remat), at 32x2x8 and 64x1x8 (512 devices;
  query_proj_kernel_gather_once off, query_proj device) and with
  TRAINER_TUNED_FLAGS=false. It refuses the 32x2x4 trainer on the previous pin
  bbd0f6f10, which does not define
  context_parallel_attention_load_balance_hoist_ids, and refuses
  query_proj_kernel_gather_once at 32x2x8 next to the block's
  query_proj=device ("query_proj_kernel_gather_once=True requires a remat
  policy that replays the query projection").

- requirements: pin MaxText to 1bf840cff, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (15285fd74 + the GDN context-parallel scan commits, their key gating and
  its test). Both new keys default to off, so the pin alone changes no
  rendered launch.
- mlperf_397b_256_v7x.sh: with TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh, the one it was timed at, the block passes
  gdn_cp_m_scan=true: with gdn_cp_mode=seq the GDN context-parallel forward
  folds each rank's chunk transitions with one Pallas scan in VMEM instead
  of an XLA reduction tree. Measured on the RL trainer bench at 32x2x4, same
  pod: 79.480 -> 78.822 s/step (3,711 -> 3,742 tokens/s/chip), 78.795 in a
  replicate in another pod, 0 tokens dropped. Its backward twin
  gdn_cp_dh_scan stays off: alone it made the step slower (79.403 -> 80.610).
- mlperf_submission_launcher_test.py: the bench cell gains exactly this key,
  once; it is absent with the tuned flags off and at the 256-chip meshes,
  and gdn_cp_dh_scan is in no render.

tests/utils/fp32_master_residual16_test.py:
- _residual16_split / _residual16_join: every non-NaN float32 pattern in a
  2^20-element random sample plus zeros, infinities, the largest finite
  value, the smallest normal, denormals and exact bfloat16 rounding ties is
  rebuilt bit for bit from its bfloat16 cast, by round-to-nearest-even and
  by truncation; the violation counter counts exactly the inputs that are
  not rebuilt, and those are all NaN.
- _pack_bits / _unpack_bits round-trip along each axis for lengths 1, 7, 8,
  9, 15, 16 and 17.
- The optimizer on a toy model (bfloat16 and float32 scanned `layers`
  weights stacked on axis 1, a non-scanned bfloat16 weight, a float32
  weight), adamw with clipping, 8 steps of random gradients: weights,
  rebuilt masters and Adam state bit-identical to the float32 storage after
  every step, residual16_violations 0; only the scanned bfloat16 weights
  store uint16 low bits and a packed uint8 bit array (stack axis 15 -> 2).
- A join that ignores the up bit is caught by the same comparison.
- The builder refuses an unknown storage, and residual16 without the scan
  axis; create_maxtext_engine passes fp32_master_storage and
  param_scan_axis to it (and builds the f32 optimizer as before).

tests/utils/maxtext_utils_test.py: the create_maxtext_engine test's
MagicMock config gains fp32_master_storage = "f32". create_maxtext_engine
now reads that key, and a MagicMock attribute is not a storage name, so
without it the test raised ValueError.

…educe-scatter flag at 32x2x4

- requirements: pin MaxText to 667955277, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (1bf840cff + gdn_cp_pass1_states_and_bf16_qkv / gdn_out_proj_token_minor,
  moe_fp8_bwd_dispatch, ragged_unsort_spread_masked_indices,
  moe_fallback_topk_capacity, fp32_master_storage, gdn_cp_fast_state_passing
  and their tests). The new keys default to off, so the pin alone changes no
  rendered launch.
- mlperf_397b_256_v7x.sh, TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh:
  - part 1: the default trainer libtpu flags carry
    --xla_tpu_use_single_sparse_core_for_reduce_scatter_offload=true
    instead of false (trainer env, Pathways worker env and proxy args; a
    caller's own TRAINER_LIBTPU_INIT_ARGS is kept);
  - part 2: gdn_cp_dh_scan=true, gdn_cp_pass1_states_and_bf16_qkv=true,
    gdn_out_proj_token_minor=true, moe_fp8_bwd_dispatch=true (only with
    TRAINER_FP8=true: it is lossy and its numerics run used the fp8 cell),
    ragged_unsort_spread_masked_indices=true, moe_fallback_topk_capacity=128,
    gdn_cp_fast_state_passing=true and fp32_master_storage=residual16.
  The comments give each key's same-pod measurement on the RL trainer bench
  (AI-Hypercomputer/maxtext#5701) and its numerics run, except
  gdn_cp_fast_state_passing, which has neither: it was not timed alone on
  TPU and has no numerics run of its own (the comments say it is
  bit-identical on CPU; on TPU it ran only inside the whole stack).
  The flag numbers come from the stacks before the final keys:
  gdn_cp_dh_scan alone was slower without the flag (+1.2 s/step, on the
  scan code without the 2048 wi tile); with the tile, the m scan and the
  pass-1 keys, the flag alone cost about 0.9 s/step and the two together
  were about 1.0 s/step faster. On the whole stack, a pod without the flag
  measured 73.200 s/step against 73.202 and 73.233 in two pods with it (a
  cross-pod comparison); the recipe keeps the flag, as the stack was timed
  with it.
- mlperf_submission_launcher_test.py: the bench cell gains exactly these
  keys and the flag; each key appears once at 32x2x4 with the tuned flags
  (moe_fp8_bwd_dispatch not with TRAINER_FP8=false) and nowhere else; the
  flag is true three times at that cell, false everywhere else, and a
  caller's own libtpu flags are kept; gdn_cp_dh_scan is no longer asserted
  absent at 32x2x4.

…tates and gdn_gated_norm_layout=token_minor at 32x2x4

- requirements: pin MaxText to ae5257c1f, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (6679552770a0 + moe_fp8_token_all_gather, gdn_bwd_natural_layout and
  gdn_gated_norm_layout and their tests). The new keys default to off, so
  the pin alone changes no rendered launch.
- mlperf_397b_256_v7x.sh, TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh: moe_fp8_token_all_gather=true (only with
  TRAINER_FP8=true: it is lossy and its numerics run used the fp8 cell),
  gdn_bwd_natural_layout=states and gdn_gated_norm_layout=token_minor.
  The comments give each key's same-pod measurement on the RL trainer bench
  (AI-Hypercomputer/maxtext#5701), on the stack of the keys before them:
  moe_fp8_token_all_gather -0.510 and -0.418 s/step in two pods, numerics
  6 of 6; gdn_bwd_natural_layout=states -0.208 (with a dead-tail-row key
  this recipe does not set, -0.025 on its own), numerics 6 of 6, and
  -0.286 on top of moe_fp8_token_all_gather; gdn_gated_norm_layout=
  token_minor -0.742 on top of both. The three together have no numerics
  run yet.
- mlperf_submission_launcher_test.py: the bench cell gains exactly these
  three keys; each appears once at 32x2x4 with the tuned flags
  (moe_fp8_token_all_gather not with TRAINER_FP8=false) and nowhere else.

…adients

With TUNIX_LOGPS_SHARDED_HEAD on, the head's backward adds every chunk's
parameter gradient into a full-size float32 accumulator per device and
reduce-scatters it once, after the chunk loop. For Qwen3.5-397B's
[4096, 248320] vocabulary table that accumulator is 3.8 GiB per device, and
XLA:TPU keeps it live across the decoder's backward and lowers the
reduce-scatter as an all-reduce plus a dynamic slice.

TUNIX_LOGPS_HEAD_GRAD_SCATTER (default off: the backward is unchanged)
reduce-scatters each chunk's gradient inside the chunk loop and keeps only
this device's float32 shard:
  chunk         reduce-scatter in the dtype the head's vjp returns (bf16 for
                MaxText's table), then add the shard in float32
  chunk_f32     cast to float32 first: the same float32 sums as off, in
                another order, at twice the wire bytes
  chunk_unroll  chunk, with the chunk loop unrolled
Each reduce-scattered shard passes an optimization barrier before it is
added: without it, XLA's while-loop all-reduce code motion moves the
reduce-scatter back after the loop, which restores the full-size
accumulator (and, with the bf16-to-f32 cast, fails a RET_CHECK in
while_loop_all_reduce_code_motion.cc in a deviceless TPU compile).
An unknown value, or the switch on while the sharded head is off, raises
ValueError. Each trace with the switch on logs one
"[logps-head] table-gradient scatter <mode>: ..." warning. Both backward
loops (the fused one-pass backward and the autodiff one) take it.

tests/rl/sharded_logps_head_test.py: parsing and refusals, parity with off
for f32 and bf16 heads in all three backward configurations, where the
reduce-scatters sit (inside the chunk loop, unrolled for chunk_unroll), the
banner and the fallback warning.

…_chunk_swap_local_self and TUNIX_LOGPS_HEAD_GRAD_SCATTER=chunk at 32x2x4

- requirements: pin MaxText to 8c56959ae, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (ae5257c1f + context_parallel_attention_load_balance_order and its test).
  The new key defaults to off, so the pin alone changes no rendered launch.
- mlperf_397b_256_v7x.sh, TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh:
  - context_parallel_attention_load_balance_order=dual_chunk_swap_local_self
    (MaxText): the attention layers' re-order into and out of the
    load-balanced token order leaves the self pairs out of its two
    collective-permutes. Bit-identical on CPU.
  - TUNIX_LOGPS_HEAD_GRAD_SCATTER=chunk (tunix/rl/common.py, the previous
    commit): the log-prob head's backward reduce-scatters each chunk's
    table gradient in bfloat16 inside the chunk loop instead of one
    float32 reduce-scatter after it. Precision class: the gradient changes
    in its low bits. It is an environment variable of the trainer process,
    so the recipe appends it to TRAINER_EXTRA_ENV, only with
    TUNIX_LOGPS_SHARDED_HEAD=fused (tunix refuses it with the sharded head
    off); a value the caller set is kept.
  The comments give the same-pod measurements on the RL trainer bench
  (AI-Hypercomputer/maxtext#5701): the re-order -0.258 s/step and the head
  reduce-scatter -0.281 s/step, each on the stack before the previous
  commit's three keys; the two together -0.542 s/step on top of all the
  keys (71.771 -> 71.229, 4,140 tokens/s/chip; 71.025 in a second pod),
  numerics 6 of 6. They also record the previous three keys' numerics run
  (6 of 6) and replicate pod, which the comments said were not run yet.
- mlperf_submission_launcher_test.py: the bench cell gains exactly the new
  key, which appears once at 32x2x4 with the tuned flags and nowhere else;
  the head variable appears once in that trainer's env (with the fused
  head only, a caller's value kept) and in no other JobSet.

…up_size=4 and gdn_in_proj_kernel_order at 32x2x4

- requirements: pin MaxText to d7c75f205, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (8c56959ae + gdn_conv1d_natural_layout / gdn_fwd_kernel_head_group_size,
  gdn_in_proj_kernel_order and a test commit). The new keys default to off,
  so the pin alone changes no rendered launch.
- mlperf_397b_256_v7x.sh, TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh, right after the attention re-order key (the order the
  TPU runs appended them in):
  - gdn_conv1d_natural_layout=true with gdn_fwd_kernel_head_group_size=4: the
    prefill-only GDN forward kernels run their Conv1D + SiLU on the bfloat16
    chunk they DMA, with tokens on sublanes, instead of through a per-token
    float32 scratch, and take 4 kq heads per group. Bit-identical on CPU.
    MaxText refuses the first without gdn_cp_pass1_states_and_bf16_qkv=true,
    which the block already sets.
  - gdn_in_proj_kernel_order=true: the GDN in_proj_qkvz is computed as two
    matmuls already in the kernel's q|k|v order (and z), from the permuted
    kernel gathered once, instead of a per-head split and relayout. Forward
    values and parameter gradients bit-identical on CPU, the input gradient
    to float32 summation order. MaxText refuses it with head-sharded GDN
    context parallelism; the block sets gdn_cp_mode=seq.
  The comments give the same-pod measurements on the RL trainer bench
  (AI-Hypercomputer/maxtext#5701), each on top of all of the keys before it
  against the same code without it: 71.180 -> 69.707 s/step for the first
  pair (4,143 -> 4,231 tokens/s/chip) and 71.180 -> 69.979 for the second
  (4,214), numerics 6 of 6 each. The two together have not run on TPU.
- mlperf_submission_launcher_test.py: the three keys in the tuned-key tables
  (default, tuned and final values, and the keys an older MaxText lacks),
  and a new test: at 32x2x4 they follow the re-order key in that order, with
  the prerequisites MaxText checks for them in the same flags.

…ident_gib=6 and ragged_buffer_factor=2.0 at 32x2x4

- requirements: pin MaxText to 1667e6dee, the new head of
  AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext
  (d7c75f205 + dense_grad_reduce, optimizer_device_resident_gib and a test
  commit). The new keys default to off, so the pin alone changes no rendered
  launch.
- mlperf_397b_256_v7x.sh, TRAINER_TUNED_FLAGS=true, only at the 128-chip
  32x2x4 trainer mesh, right after the GDN keys (the order the TPU runs
  appended them in):
  - dense_grad_reduce=all_reduce: the dense-weight gradients that
    cp-as-ep-dense-no-cp replicates over context are each reduced by one
    all-reduce instead of a reduce-scatter over all devices and an all-gather
    over context. The forward is unchanged; the gradients differ in summation
    order only. MaxText refuses it under a rule set without embed_grad; the
    block sets cp-as-ep-dense-no-cp.
  - optimizer_device_resident_gib=6, on its own line: the largest optimizer
    leaves, up to 6 GiB per device, stay in device memory between updates,
    read and written in place, instead of crossing the host link twice per
    step. Placement only. MaxText refuses it without
    optimizer_memory_host_offload=true, which the recipe sets.
  - ragged_buffer_factor=2.0, only with RETRY_WHEN_TOKENS_DROPPED=true and a
    finite RAGGED_BUFFER_FACTOR: a micro-batch whose routing overflows the
    smaller buffer is rerun with the dropless program, so no token is
    dropped. RAGGED_BUFFER_FACTOR=-1 and launches without the retry keep
    their own value.
  The comments give the same-pod measurements on the RL trainer bench
  (AI-Hypercomputer/maxtext#5701): each of the first two on top of the keys
  before it, against the same code without it (71.180 -> 69.446 and -> 70.576
  s/step), and the two with the two GDN keys on the bench's production
  routing table (69.339 -> 64.271 s/step, 4,253 -> 4,589 tokens/s/chip),
  numerics 6 of 6.
- mlperf_submission_launcher_test.py: the two keys in the tuned-key tables
  (default, tuned and final values, and the keys an older MaxText lacks), the
  buffer factor the tuned block sets at 32x2x4 with the retry, and a new test:
  at 32x2x4 the three follow the GDN keys in that order, with the
  prerequisites MaxText checks for the first two in the same flags, and a
  fully dropless or no-retry launch keeps its own buffer factor.